fix(core): Overhaul worker shutdown to make sure all tasks are executed before terminating - #1704

I think we can simplify this a fair amount. Right now it's easy to forget to use the semaphore properly and we could solve that by reversing the way it works. EX: If anyone forgets to pass in an attempt semaphore, that call can get interrupted immediately.

Rather than initializing it with a permit, we should init it with 0, then the permit gets added in TonicErrorHandler::wait_retry and every attempt after the first attempt acquires in the gRPC interceptor but wait to forget it until just before running.

Here's an AI sketch of the changes:

1. crates/client/src/grpc.rs: only retries acquire a permit

let maybe_sem = req.extensions().get::<AttemptSemaphore>();
let mut is_first_attempt = true;

let fact = || {
    let req_clone = req_cloner(&req);
    let call_fut = callfn(req_clone);
    // The permit for a retry is the one TonicErrorHandler added when the previous attempt failed.
    let gate = (!std::mem::take(&mut is_first_attempt)).then_some(maybe_sem).flatten();
    async move {
        if let Some(sem) = gate {
            sem.0.acquire().await.expect("AttemptSemaphore must not be closed").forget();
        }
        call_fut.await
    }
};

make_future_retry takes an FnMut, so mutating the flag inside the closure is fine.

2. crates/client/src/retry.rs: no change

wait_retry already does sem.add_permits(1) before returning WaitRetry. The ForwardError and success paths leave the count at 0, as they do now.

3. crates/client/src/request_extensions.rs: update the doc

/// A request extension that exposes when a call is between retry attempts. The semaphore must start
/// with 0 permits. A permit is added after an attempt fails with a retryable error, and taken back
/// before the next attempt starts. Acquiring a permit from outside therefore succeeds only between
/// attempts, which is the only point at which the call can be abandoned without losing a response.
#[derive(Clone, Debug)]
pub struct AttemptSemaphore(pub Arc<Semaphore>);

4. crates/sdk-core/src/pollers/poll_buffer.rs: start at 0, poll branch first

let sem = Arc::new(Semaphore::new(0));
(Some(sem.clone()), async move {
    // Only resolves between retry attempts, so no server response can be dropped.
    if let Ok(permit) = sem.acquire().await {
        permit.forget();
    }
}.left_future())
...
let r = tokio::select! {
    biased;
    r = pf(PollFnOptions { timeout_override, attempt_semaphore }) => r,
    _ = async { shutdown.cancelled().await; poll_interruptor.await } => return,
};

Putting the poll branch first also fixes another issue: a poll that has just completed wins the tie. That's safe in graceful mode, because the waiting interrupt branch is first in the semaphore's FIFO queue. When a retry adds a permit, it goes to the interrupt branch, not to the next attempt.

5. Mocks and tests

• Remove both "MUST lock the semaphore" notes in mocks.rs. Mocks that ignore the semaphore now behave like the old code: polls run to completion.
• shutdown_blocks_until_shutdown_rpc_first_attempt_done: the mock no longer needs to touch the semaphore.
• shutdown_proceeds_on_poll_retryable_error: simulate a retry with sem.add_permits(1) instead of try_acquire/drop.
• Add one unit test in grpc.rs: the first attempt goes through with 0 permits, and the second waits until add_permits(1).

The other thing we can do I think is eliminate RetryCallback entirely, and re-use the semaphore to deal with the same thing:

// initiate_shutdown
let between_attempts = Arc::new(Semaphore::new(0));
let join_handle = tokio::spawn(async move {
    client.shutdown_worker(sticky_name, task_queue, task_queue_types, heartbeat,
                           Some(between_attempts_clone)).await ...
});
*guard = Some(ShutdownRpcContext { join_handle, between_attempts });

// shutdown
select! {
    res = &mut ctx.join_handle => handle_join(res),
    _ = ctx.between_attempts.acquire() => {
        warn!("ShutdownWorker RPC still failing after pollers stopped; abandoning it");
        ctx.join_handle.abort();
        #[cfg(debug_assertions)]
        handle_join(ctx.join_handle.await);
    }
}

It does mean we will now possibly wait for an in-flight ShutdownWorker rpc to finish, but, that's totally fine.

Same case twice

Leftover code

Leftover code

I think there's still a problem here. Since PollerStopped gets passed straight through this, regardless of if the eager side is done, which means the poller gets marked stopped and can finish shutdown & drop external_wft_rx. Worth adding a regression test that tries to trigger an eager task getting dropped.

Can be fixed I think by doing the same thing we do for sticky poller: https://github.com/temporalio/sdk-rust/pull/1704/changes#diff-e5cb102b5d9c13362ffe79169f2f125a11add40b80095e1d8cf2221ec1fdd998R81

Perhaps we want to wrap that up in some shared helper

I think here we could just call state.poller.shutdown().await directly, which would allow us to just get rid of the oneshot channels entirely.

This would make StickyTaskPoller::shutdown not behave correctly, and downstream callers rely on it working.

This change makes the comment wrong now... why did we go from 1 -> None?

Copy-paste error. Will fix.

Why did this assertion go away?

Not sure. But it seems to work fine with it added back, so I'll add it back.